import os

from dotenv import load_dotenv

from lc_frame.lc_models.api.llm_api import get_qwen
from lc_frame.lc_agent.agent import lc_agent
from lc_frame.lc_rag.lc_vector import lc_rag
from lc_frame.utils.log import logger
from langchain_core.prompts import FewShotChatMessagePromptTemplate, ChatPromptTemplate, SystemMessagePromptTemplate, \
    HumanMessagePromptTemplate
from langchain_core.output_parsers import StrOutputParser

current_path = os.path.abspath(__file__)
root_path = os.path.dirname(os.path.dirname(current_path))

# 加载key
dotenv_path = os.path.join(root_path, "lc_models", "api", ".qwen")
load_dotenv(dotenv_path=dotenv_path)

chat, llm, embed = get_qwen()


def recognition_intent(query):
    """
        基于 LLM，进行少量样本学习，对用户的输入文本进行意图识别，是进行RAG检索还是agent进行nl2sql
    """
    # 进行少量样板学习，模版参照官网样例
    # https://python.langchain.com/api_reference/core/prompts/langchain_core.prompts.few_shot.FewShotChatMessagePromptTemplate
    # 带有具体日期的问题，进行数据库查询，没有的进行文档查询
    examples = [
        {"input": "国元证券拟上市证券交易所是在哪里？", "output": "RAG|国元证券拟上市证券交易所是在哪里？"},
        {"input": "青海明胶上市时间是哪年？", "output": "RAG|青海明胶上市时间是哪年？"},
        {"input": "安徽黄山胶囊股份有限公司法人是谁？", "output": "RAG|安徽黄山胶囊股份有限公司法人是谁？"},
        {"input": "国元证券股份有限公司住所是哪里？", "output": "RAG|国元证券股份有限公司住所是哪里？"},
        {"input": "国元证券股份有限公司的董事长是谁？", "output": "RAG|国元证券股份有限公司的董事长是谁？"},

        {"input": "2019年，深圳证券交易所下有多少股票？", "output": "Agent|2019年，深圳证券交易所下有多少股票？"},
        {"input": "20190109日，股票代码[2280 HK]的成交金额是多少？",
         "output": "Agent|20190109日，股票代码[2280 HK]的成交金额是多少？"},
        {"input": "2019年农银汇理金汇债券A的公告日期是哪天？",
         "output": "Agent|2019年农银汇理金汇债券A的公告日期是哪天？"},
        {"input": "2019年中信保诚嘉鸿债券C的报告类型是什么？",
         "output": "Agent|2019年中信保诚嘉鸿债券C的报告类型是什么？"},
        {"input": "20190131日，股票代码[603665]的成交金额是多少？",
         "output": "Agent|20190131日，股票代码[603665]的成交金额是多少？"},
        {"input": "圳信立泰药业股份有限公司注射用头孢西丁钠（信希汀、1.0克装）在2006 年、2007 年、2008 年、2009 年1-6 月的最高零售价格分别为多少",
         "output": "RAG|圳信立泰药业股份有限公司注射用头孢西丁钠（信希汀、1.0克装）在2006 年、2007 年、2008 年、2009 年1-6 月的最高零售价格分别为多少"},
        {"input": "2005年、2006年和2007年，宁波立立电子股份有限公司所获政府补助计入当期损益的金额分别占当年利润总额的比例是多少？",
         "output": "RAG|2005年、2006年和2007年，宁波立立电子股份有限公司所获政府补助计入当期损益的金额分别占当年利润总额的比例是多少？"},
        {"input": "我想知道股票600508在申万行业分类下的二级行业是什么？用最新的数据。",
         "output": "Agent|我想知道股票600508在申万行业分类下的二级行业是什么？用最新的数据。"}
        
        
    ]

    example_prompt = ChatPromptTemplate.from_messages(
        [('human', '{input}'), ('ai', '{output}')]
    )

    few_shot_prompt = FewShotChatMessagePromptTemplate(
        examples=examples,
        example_prompt=example_prompt,
    )

    final_prompt = ChatPromptTemplate.from_messages(
        [
            ('system', """
            请学习我给定的样例，并据此回答我提出的问题:
            如果既不属于Agent也不属于RAG,就回答other;
            你只能回答Agent|原始输入、RAG|原始输入或者other|原始输入，
            例如: other|今天天气真好:\n
            """),
            few_shot_prompt,
            ('human', '{input}'),
        ]
    )

    # 这里只是一个示例，实际上你会调用一个语言模型
    chain = final_prompt | chat | StrOutputParser()
    response = chain.invoke({"input": query})
    logger.info(f"意图识别后返回的结果为：{response}")
    return response


def do_action(query):
    """
        基于 LLM，根据意图和查询，返回相应的结果
    """
    intent = query.split("|")[0]
    question = query.split("|")[1]
    logger.info(f"意图识别后为：{intent}")
    result = ""
    if intent == "RAG":
        result = lc_rag(question=question)

    elif intent == "Agent":
        result = lc_agent(question=question)
    else:
        #result = chat.invoke(input=question)
        result = chat.invoke(input=question).content
    return result,intent


def get_final_result(query, result):
    """
    融合信息
    """
    messages = [
        SystemMessagePromptTemplate.from_template(template="你是一个信息融合机器人"),
        HumanMessagePromptTemplate.from_template("""
                                                 请将{role}中的信息进行融合，其中输入的信息为2部分：
                                                 用户输入的问题/回答结果，
                                                 请将输入的信息融合，用更加合理通顺的语句将信息整合在一起，
                                                 直接输出融合后得结果即可，不需要输出原始的回答
                                                 """),
    ]

    prompt = ChatPromptTemplate.from_messages(messages=messages)

    chain = prompt | chat

    final_result = chain.invoke(input={"role": query + result})
    #return final_result
    return final_result.content


def intent_handler(query):
    """
    主函数，进行意图识别和查询
    """
    intent_result = recognition_intent(query=query)
    result,intent = do_action(intent_result)
    final_result = get_final_result(query, result)
    return final_result,intent,result=="不知道！"


# if __name__ == '__main__':
#     query = "农银汇理金汇债券A的公告日期是哪天？"
#     final_result,_ = intent_handler(query)
#     logger.info(f"最终结果为：{final_result}")
